import sys
import re
import os
import subprocess
import shutil
import argparse

def extract_fortran_elements(file_path):
    """Extract subroutines, modules, and programs from a Fortran file."""
    subroutines = []
    modules = []
    programs = []
    current_element = None
    current_content = []
    element_type = None
    
    with open(file_path, 'r') as f:
        lines = f.readlines()
        for i, line in enumerate(lines):
            line_stripped = line.strip()
            if line_stripped.startswith('!') or not line_stripped:
                # Skip comments and empty lines
                continue
            
            # Check for MODULE
            module_match = re.match(r'(?i)^\s*MODULE\s+(\w+)', line_stripped)
            if module_match:
                if current_element:
                    if element_type == 'subroutine':
                        subroutines.append((current_element, current_content[:]))
                    elif element_type == 'module':
                        modules.append((current_element, current_content[:]))
                    elif element_type == 'program':
                        programs.append((current_element, current_content[:]))
                current_element = module_match.group(1)
                current_content = [line]  # Store the original line with whitespace
                element_type = 'module'
                continue
            
            # Check for SUBROUTINE
            subroutine_match = re.match(r'(?i)^\s*SUBROUTINE\s+(\w+)', line_stripped)
            if subroutine_match:
                if current_element:
                    if element_type == 'subroutine':
                        subroutines.append((current_element, current_content[:]))
                    elif element_type == 'module':
                        modules.append((current_element, current_content[:]))
                    elif element_type == 'program':
                        programs.append((current_element, current_content[:]))
                current_element = subroutine_match.group(1)
                current_content = [line]  # Store the original line with whitespace
                element_type = 'subroutine'
                continue
            
            # Check for PROGRAM
            program_match = re.match(r'(?i)^\s*PROGRAM\s+(\w+)', line_stripped)
            if program_match:
                if current_element:
                    if element_type == 'subroutine':
                        subroutines.append((current_element, current_content[:]))
                    elif element_type == 'module':
                        modules.append((current_element, current_content[:]))
                    elif element_type == 'program':
                        programs.append((current_element, current_content[:]))
                current_element = program_match.group(1)
                current_content = [line]  # Store the original line with whitespace
                element_type = 'program'
                continue
            
            # Check for END statements
            end_match = re.match(r'(?i)^\s*END\s+(MODULE|SUBROUTINE|PROGRAM)?\s*(\w*)', line_stripped)
            if end_match and current_element:
                current_content.append(line)  # Store the original line with whitespace
                if element_type == 'subroutine':
                    subroutines.append((current_element, current_content[:]))
                elif element_type == 'module':
                    modules.append((current_element, current_content[:]))
                elif element_type == 'program':
                    programs.append((current_element, current_content[:]))
                current_element = None
                current_content = []
                element_type = None
                continue
            
            if current_element:
                current_content.append(line)  # Store the original line with whitespace
    
    # Handle case where file ends without an END statement
    if current_element:
        if element_type == 'subroutine':
            subroutines.append((current_element, current_content[:]))
        elif element_type == 'module':
            modules.append((current_element, current_content[:]))
        elif element_type == 'program':
            programs.append((current_element, current_content[:]))
    
    return subroutines, modules, programs

def list_elements(subroutines, modules, programs):
    """List all subroutines, modules, and programs."""
    print("Subroutines:")
    for name, _ in subroutines:
        print(f"  - {name}")
    print("Modules:")
    for name, _ in modules:
        print(f"  - {name}")
    print("Programs:")
    for name, _ in programs:
        print(f"  - {name}")

def print_function_context(func_name, subroutines, modules, programs):
    """Print the context of a specific function, module, or program."""
    for name, content in subroutines:
        if name.lower() == func_name.lower():
            print(f"Context of SUBROUTINE {name}:")
            print(''.join(content))
            return content
    for name, content in modules:
        if name.lower() == func_name.lower():
            print(f"Context of MODULE {name}:")
            print(''.join(content))
            return content
    for name, content in programs:
        if name.lower() == func_name.lower():
            print(f"Context of PROGRAM {name}:")
            print(''.join(content))
            return content
    print(f"Function {func_name} not found.")
    return None

def translate_to_c(func_content, func_name):
    """Translate Fortran content to C using autoF2C.py."""
    # Create clean working directory
    if os.path.exists('output'):
        # Clean up existing directory but keep error log
        for file in os.listdir('output'):
            if file != 'error_log' and os.path.isfile(os.path.join('output', file)):
                os.remove(os.path.join('output', file))
            if os.path.isdir(os.path.join('output', file)) and file != 'c':
                shutil.rmtree(os.path.join('output', file))
    else:
        os.makedirs('output')
    
    # Delete existing c directory if it exists
    if os.path.exists('c'):
        shutil.rmtree('c')
    
    # Create a clean, properly formatted Fortran file
    # This is a modified version of the original code to ensure proper format
    fortran_code = """SUBROUTINE solve_helmholts(a_helm,b_helm,cm_helm,threshold, pi, &
                       idep, jdep,         &
                       ids,ide,jds,jde,kds,kde,   &
                       ims,ime,jms,jme,kms,kme,   &
                       its,ite,jts,jte,kts,kte,myprcid)

  IMPLICIT NONE

  INTEGER   ::           ims,ime,jms,jme,kms,kme,   &
                         ids,ide,jds,jde,kds,kde,   &
                         its,ite,jts,jte,kts,kte,   &
                         idep, jdep

  REAL*8,DIMENSION(ims:ime,kms:kme,jms:jme),INTENT(inout) :: pi
  REAL*8,DIMENSION(ims:ime,kms:kme,jms:jme):: x
!local
  REAL*4                               :: ep, threshold
  REAL*8                               :: d
  INTEGER,parameter                              :: iter_max=10
  INTEGER                              :: int1,int2
  INTEGER                              :: i,j,k,n,l,kk, myprcid
  INTEGER                              ::  ibegin,iend,jbegin,jend
  REAL,   DIMENSION(1:19,its:ite,kts-1:kte+1,jts:jte) ::   a_helm
  REAL,   DIMENSION(its:ite,kts-1:kte+1,jts:jte)  ::   b_helm
  REAL,   DIMENSION(ims:ime,kms:kme,jms:jme,1:7) ::   cm_helm
#ifdef GPTL_TIMING
   integer   :: ret
    ret=GPTLsetoption (GPTLoverhead, 0); 
   ret = gptlinitialize()
   ret = gptlstart('gcr_main')
#endif

  ibegin = its
  iend   = ite
  jbegin = jts
  jend   = min (jde-1,jte)


  d=1.0d0
  ep=threshold
DO i=1,1
  CALL psolve_gcr_main(ep,a_helm,b_helm,cm_helm,iter_max,pi,d, &
                       idep, jdep,         &
                       ids,ide,jds,jde,kds,kde,   &
                       ims,ime,jms,jme,kms,kme,   &
                       its,ite,jts,jte,kts,kte,myprcid)
enddo
#ifdef GPTL_TIMING
   ret = gptlstop('gcr_main')
   ret = gptlpr_file("outfile_f90")
#endif


#ifdef GPTL_TIMING
   ret = gptlfinalize()
#endif

END SUBROUTINE solve_helmholts
"""
    
    # Use a predictable filename in a clean directory
    temp_file = os.path.join('output', f'solve_helmholts.F90')
    
    with open(temp_file, 'w') as f:
        f.write(fortran_code)
    
    print(f"Exported formatted Fortran code to {temp_file}")
    
    # Run autoF2C_fixed.py (the fixed version) with proper parameters
    try:
        print(f"Running: python ./autoF2C.py {temp_file}")
        
        # Capture stdout and stderr for better error reporting
        result = subprocess.run(['python', './autoF2C.py', temp_file], 
                                capture_output=True, text=True)
        
        print(result.stdout)  # Print the output for debugging
        
        if result.returncode == 0:
            print("autoF2C_fixed.py execution completed")
            
            # Look for the C file in the c/ directory (from the root)
            c_file_path = os.path.join('c', 'solve_helmholts.c')
            
            if os.path.exists(c_file_path):
                print(f"Found C output at: {c_file_path}")
                
                # Verify the file contains more than just headers
                with open(c_file_path, 'r') as f:
                    content = f.read()
                    if len(content.strip()) > 50:  # Arbitrary length check
                        print(content)
                    else:
                        print(f"Warning: C file exists but appears incomplete: {content}")
            else:
                print("C file not found in the expected location: c/solve_helmholts.c")
                
                # Search for any .c files as a fallback
                print("Searching for any .c files...")
                for root, dirs, files in os.walk('.'):
                    for file in files:
                        if file.endswith('.c'):
                            c_path = os.path.join(root, file)
                            print(f"Found C file: {c_path}")
                            with open(c_path, 'r') as f:
                                content = f.read()
                                if len(content.strip()) > 50:
                                    print(content)
                                else:
                                    print(f"Warning: C file exists but appears incomplete: {content}")
        else:
            print("Translation failed:")
            print(result.stderr)
            
            # Check error log
            error_log_path = os.path.join('output', 'error_log')
            if os.path.exists(error_log_path):
                print("\nContent of error_log:")
                with open(error_log_path, 'r') as f:
                    print(f.read())
    
    except Exception as e:
        print(f"Error during translation: {str(e)}")

def extract_function(input_file, function_name):
    """Extract the specified function/subroutine from the Fortran file."""
    with open(input_file, 'r') as f:
        content = f.read()
    
    # Find the subroutine block
    start_pattern = rf'(?i)subroutine\s+{function_name}\s*\('
    end_pattern = rf'(?i)end\s+subroutine\s*{function_name}'
    start_match = re.search(start_pattern, content)
    end_match = re.search(end_pattern, content)
    
    if start_match and end_match:
        start_idx = start_match.start()
        end_idx = end_match.end()
        return content[start_idx:end_idx]
    else:
        print(f"Error: Function {function_name} not found in {input_file}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='Test Fortran to C conversion.')
    parser.add_argument('input_file', help='Input Fortran file')
    parser.add_argument('-f', '--function', help='Specific function/subroutine to convert')
    parser.add_argument('--trans', action='store_true', help='Translate the code')
    parser.add_argument('--qwen', action='store_true', help='Use Qwen model for translation')
    parser.add_argument('-l', '--list', action='store_true', help='List all subroutines, modules, and programs in the input file')
    args = parser.parse_args()

    if not os.path.exists(args.input_file):
        print(f"Error: Input file {args.input_file} not found.")
        sys.exit(1)

    if args.list:
        subroutines, modules, programs = extract_fortran_elements(args.input_file)
        list_elements(subroutines, modules, programs)
    elif args.function and args.trans:
        # Create output directory
        os.makedirs('output', exist_ok=True)
        output_file = os.path.join('output', f"{args.function}.F90")
        
        # Extract the specific function
        function_code = extract_function(args.input_file, args.function)
        
        # Write the extracted function to a file
        with open(output_file, 'w') as f:
            f.write(function_code)
        print(f"Exported formatted Fortran code to {output_file}")
        
        # Run the appropriate conversion script
        if args.qwen:
            print(f"Running: python ./autoF2C-qwen.py {output_file}")
            try:
                subprocess.run(['python', './autoF2C-qwen.py', output_file], check=True)
                print("autoF2C-qwen.py execution completed")
                
                # Check for output file
                c_file = os.path.splitext(output_file)[0] + '.c'
                if os.path.exists(c_file):
                    print(f"Found C file: {c_file}")
                    with open(c_file, 'r') as f:
                        print(f.read())
                else:
                    print(f"C file not found at expected location: {c_file}")
            except subprocess.CalledProcessError:
                print("Error occurred during conversion with autoF2C-qwen.py")
        else:
            print(f"Running: python ./autoF2C.py {output_file}")
            try:
                subprocess.run(['python', './autoF2C.py', output_file], check=True)
                print("autoF2C.py execution completed")
                
                # Check for output file
                c_file = os.path.splitext(output_file)[0] + '.c'
                if os.path.exists(c_file):
                    print(f"Found C file: {c_file}")
                    with open(c_file, 'r') as f:
                        print(f.read())
                else:
                    print(f"C file not found at expected location: {c_file}")
                    # Search for any .c files in tmp
                    print("Searching for any .c files...")
                    for root, dirs, files in os.walk('output'):
                        for file in files:
                            if file.endswith('.c'):
                                c_file_path = os.path.join(root, file)
                                print(f"Found C file: {c_file_path}")
                                with open(c_file_path, 'r') as f:
                                    print(f.read())
                                break
            except subprocess.CalledProcessError:
                print("Error occurred during conversion with autoF2C.py")
    else:
        print("Please specify a function with -f and enable translation with --trans")

if __name__ == '__main__':
    main() 